package safe

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unsafe"
)

var (
	kernel32                 = syscall.NewLazyDLL("kernel32.dll")
	procEnumProcessModules   = kernel32.NewProc("K32EnumProcessModules")
	procGetModuleFileNameExW = kernel32.NewProc("K32GetModuleFileNameExW")
	procGetCurrentProcess    = kernel32.NewProc("GetCurrentProcess")
	checkInterval            = time.Second // 定期检查间隔
)

// 定期扫描已加载的 DLL
func monitorDLLs(currentExecutable string) {
	for {
		err := CheckInjectedDLLs(currentExecutable)
		if err != nil {
			fmt.Printf("检测到非法 DLL 注入，程序即将退出: %v\n", err)
			os.Exit(1) // 立即退出程序
		}
		time.Sleep(checkInterval)
	}
}

// CheckInjectedDLLs 检查当前加载的 DLL 是否安全
func CheckInjectedDLLs(currentExecutable string) error {
	// 获取当前进程句柄
	handle, _, _ := procGetCurrentProcess.Call()
	if handle == 0 {
		return fmt.Errorf("无法获取当前进程句柄")
	}

	var modules [1024]uintptr
	var cbNeeded uint32

	// 获取当前进程加载的模块句柄
	ret, _, _ := procEnumProcessModules.Call(
		handle,
		uintptr(unsafe.Pointer(&modules)),
		uintptr(len(modules)*int(unsafe.Sizeof(uintptr(0)))),
		uintptr(unsafe.Pointer(&cbNeeded)),
	)
	if ret == 0 {
		return fmt.Errorf("无法枚举模块")
	}

	moduleCount := int(cbNeeded / uint32(unsafe.Sizeof(uintptr(0))))
	for i := 0; i < moduleCount; i++ {
		var modulePath [syscall.MAX_PATH]uint16

		// 获取模块文件路径
		ret, _, _ := procGetModuleFileNameExW.Call(
			handle,
			modules[i],
			uintptr(unsafe.Pointer(&modulePath)),
			uintptr(len(modulePath)),
		)
		if ret == 0 {
			continue
		}

		path := syscall.UTF16ToString(modulePath[:])
		absPath, _ := filepath.Abs(path) // 转换为绝对路径

		// 跳过当前主模块的路径
		if absPath == currentExecutable {
			continue
		}

		// 检查路径是否合法
		if !isSafeDLL(absPath) {
			return fmt.Errorf("发现非法 DLL: %s", absPath)
		}
	}
	return nil
}

// 判断 DLL 是否来自安全目录
func isSafeDLL(path string) bool {
	// 只允许加载 C:\Windows 路径下的 DLL
	//if strings.HasPrefix(path, "C:\\Program Files (x86)") {
	//	fmt.Println(path)
	//	return true
	//}
	return len(path) >= 10 && (strings.ToUpper(path[:10]) == `C:\WINDOWS`)
}
